#!/usr/bin/env python
"""

    download_ucsc_genomic_sequence.py
    
    Downloads specified genomes from UCSC
    E.g.
    download_ucsc_genomic_sequence.py
            --species_name      'gallus gallus'
            --ucsc_genome_name  galGal2
            --temp_directory    /tmp
            --goldenpath_ftp    'ftp://hgdownload.cse.ucsc.edu/goldenPath'
            --output_dir        '/net/cpp-mirror/databases/indexed_fasta'

"""

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   constants


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
UCSC_GOLDEN_PATH_FTP        = "ftp://hgdownload.cse.ucsc.edu/goldenPath"
TMP_DIR                     = "/tmp"


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   options        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import sys, os
import os.path
import StringIO
import optparse

# add self to search path for testing
if __name__ == '__main__':
    exe_path = os.path.split(os.path.abspath(sys.argv[0]))[0]
    python_modules_path = os.path.abspath(os.path.join(exe_path,"..", "python_modules"))
    sys.path.append(python_modules_path)
    module_name = os.path.split(sys.argv[0])[1]
    module_name = os.path.splitext(module_name)[0];
else:
    module_name = __name__



if __name__ == '__main__':
    parser = optparse.OptionParser(version="%prog 1.0", usage = "\n\n    %progs [options]")
    parser.add_option("-s", "--species_name", dest="species_name",
                        metavar="SPECIES", 
                        type="string",
                      help="2 part Latin name of species to download.")
    parser.add_option("-u", "--ucsc_genome_name", dest="ucsc_genome_name",
                        metavar="NAME", 
                        type="string",
                      help="Name and verion of UCSC genomic sequences to download.")
    parser.add_option("-T", "--temp_directory", dest="temp_directory",
                      metavar="PATH", 
                      default=TMP_DIR,
                      type="string",
                      help="Temporary directory for downloading and unzipping etc.")
    parser.add_option("--goldenpath_ftp", dest="goldenpath_ftp",
                      metavar="FTP_PATH", 
                      type="string",
                      default = UCSC_GOLDEN_PATH_FTP,
                      help="FTP address of goldenpath files at UCSC.")
    parser.add_option("--output_dir", dest="output_dir",
                      metavar="PATH", 
                      type="string",
                      help="Path to save downloaded FASTA files and associated indices")
    
    #
    #   general options: verbosity / logging
    # 
    parser.add_option("-v", "--verbose", dest = "verbose",
                      action="count", default=0,
                      help="Print more verbose messages for each additional verbose level.")
    parser.add_option("-L", "--log_file", dest="log_file",
                      metavar="FILE", 
                      type="string",
                      help="name and path of log file")
    parser.add_option("--skip_parameter_logging", dest="skip_parameter_logging",
                      action="store_true", default=False,
                      help="Do not print program parameters to log.")
    parser.add_option("--debug", dest="debug",
                      action="store_true", default=False,
                      help="Set default program parameters in debugging mode.")
    
    
    
    
    
    # get help string
    f =StringIO.StringIO()
    parser.print_help(f)
    helpstr = f.getvalue()
    (options, remaining_args) = parser.parse_args()
    
    if not options.log_file:
        options.log_file = os.path.join(options.output_dir, "download_ucsc_genomes.log")

    if options.debug:
        options.verbose = 1
        options.ucsc_genome_name = "galGal2"
        options.species_name = "gallus gallus"
        options.output_dir     = "/net/cpp-mirror/databases/indexed_fasta"
        options.temp_directory = os.path.join(options.output_dir, "temp")
        #options.goldenpath_ftp = "..."
    
    # mandatory options
    from options import check_mandatory_options
    mandatory_options = ["ucsc_genome_name", "species_name", "temp_directory", "goldenpath_ftp", "output_dir"]
    check_mandatory_options (options, mandatory_options, helpstr)

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   imports        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import StringIO
import re
import operator
import json
import traceback
import shutil

from collections import defaultdict
import logging
from lg_program_logging import setup_std_logging
from options import get_option_strings
import sys, os
import urllib2
import glob
from fasta_sequences_iterator import fasta_sequences_iterator
from fasta_single_line_sequence import fasta_one_line_per_sequence
from index_fasta import write_fasta_index_to_file
from get_genomic_sequence import get_genomic_sequence_coordinates_from_index_file,get_genomic_sequence



from exceptions import Exception



#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Functions        


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888
def __initial_lowercase (s):
    """
    useful for UCSC species names
    """
    return s[0:1].lower() + s[1:]



# _________________________________________________________________________________________
#
#   download_and_index_ucsc_genome_sequence
#
#_________________________________________________________________________________________
def download_and_index_ucsc_genome_sequence(logger, 
                                            species_name,
                                            ucsc_genome_name,                            
                                            genomic_sequence_dir,
                                            goldenpath_ftp = UCSC_GOLDEN_PATH_FTP, 
                                            tmp_dir = TMP_DIR,
                                            verbose = 0):
    """
    Download, unzip, index genomic sequences from ucsc
    """
    cleanup_files = set()
    try:
        os.makedirs(tmp_dir)
    except OSError, e:
        if "File exists" not in e:
            raise
        
    
    #
    #   1) define genome sequence and indices file names
    # 
    fasta_file_root  = get_fasta_name_by_convention(species_name, ucsc_genome_name, 
                                                        genomic_sequence_dir)
    fasta_file_name  = fasta_file_root + ".fasta"
    fasta_index_name = fasta_file_root + ".idx"

    
    #
    #   2) download and unzip
    # 
    unzipped_name = os.path.join(tmp_dir, ucsc_genome_name + ".fa")
    cleanup_files.add(unzipped_name)
    if os.path.exists(fasta_file_name):
        logger.info("%s downloaded" % fasta_file_name)
        logger.debug("%s downloaded" % fasta_file_name)
    else:
        if verbose > 1:
            logger.debug("Need to download %s" % fasta_file_name)
        (action, 
         local_file_name) = _download_ucsc_genomes(logger, 
                                                  ucsc_genome_name,
                                                  goldenpath_ftp,
                                                  tmp_dir,
                                                  cleanup_files,
                                                  verbose
                                                )  


        __unzip_files(logger, action, local_file_name, unzipped_name)
        

    #
    #   3) One line per contig
    #                 
    __one_line_per_contig(logger, unzipped_name, fasta_file_name)

    #
    #   4) index fasta
    #                 
    if not os.path.exists(fasta_file_name):
        raise Exception("Missing %s" % fasta_file_name)
    if not os.path.exists(fasta_index_name):
        logger.debug("Indexing %s" % fasta_file_name)
        write_fasta_index_to_file (fasta_file_name, fasta_index_name, logger)
        logger.info("%s indexed to %s" % (fasta_file_name, fasta_index_name))
    else:
        logger.debug("%s already indexed" % fasta_file_name)
        logger.info("%s already indexed" % fasta_file_name)
            

    #   cleanup if necessary 
    #
    for f in cleanup_files:
        if os.path.exists(f):
            os.remove(f)

#_________________________________________________________________________________________
#
#   1) get_fasta_name_by_convention
#_________________________________________________________________________________________
def get_fasta_name_by_convention (species_name, ucsc_genome_name, genomic_sequence_dir):
    """
    This is how Andreas Heger constructs his fasta file names.
    Kept for historical continuity with the rest of the lab...
    """
    species_names = species_name.replace(' ', '_').split('_', 1)
    if len(species_names) != 2:
        raise Exception(("Can't construct canonical species name for %s: " % species_name) +
                        "Ask Andreas how to handle this!")
    species_initials = (species_names[0][0] + species_names[1][0]).lower()
    return os.path.join(genomic_sequence_dir,
                         "%s_%s_softmasked" % (species_initials, 
                                               __initial_lowercase(ucsc_genome_name)))



#_________________________________________________________________________________________
#
#   2) _download_ucsc_genomes
#
#_________________________________________________________________________________________
def _download_ucsc_genomes(logger, 
                          ucsc_genome_name,
                          goldenpath_ftp,
                          tmp_dir,
                          cleanup_files,
                          verbose = 0):
    """
    Download and index all genomic sequence files if they are not present
    """

        
    # try all variants of the file name
    # UCSC seems to be a bit inconsistent
    remote_file_names = []
    error_messages = ""
    remote_path = goldenpath_ftp + "/%s/bigZips/" % __initial_lowercase(ucsc_genome_name)
    local_path  = tmp_dir
    remote_file_names.append(["tar", 
                              remote_path + "chromFa.tar.gz",
                              os.path.join(local_path, ucsc_genome_name + ".tar.gz")])
    remote_file_names.append(["gunzip", 
                              remote_path + "%s.fa.gz" % __initial_lowercase(ucsc_genome_name),
                              os.path.join(local_path, ucsc_genome_name + ".gz")])
    remote_file_names.append(["unzip", 
                              remote_path + "chromFa.zip",
                              os.path.join(local_path, ucsc_genome_name + ".zip")])
    remote_file_names.append(["unzip", 
                              remote_path + "ScaffoldFaMasked.zip",
                              os.path.join(local_path, ucsc_genome_name + ".zip")])

    #
    # check if exists
    # 
    for action, remote_file_name, local_file_name in remote_file_names:
        if os.path.exists(local_file_name):
            logger.debug("%s already downloaded" % local_file_name)
            cleanup_files.add(local_file_name)
            return (action, local_file_name)

    #
    #   Download if necessary:
    # 
    #
    #   try each of the download names in turn
    #
    CNT_RETRIES = 5
    for action, remote_file_name, local_file_name in remote_file_names:
        for tries in range(CNT_RETRIES):
            if verbose:
                sys.stderr.write(".")
            try:
                # does remote file exist?
                f = urllib2.urlopen(remote_file_name)
                
                # yes: download
                cleanup_files.add(local_file_name + ".tmp")
                logger.debug("Downloading %s..." % remote_file_name)
                logger.debug(" to %s..." % local_file_name)
                outfile = open(local_file_name + ".tmp", "wb")
                while True:
                    chunk = f.read(64000)
                    if not chunk:
                        break
                    outfile.write(chunk)
                outfile.close()
                #urllib2.urlretrieve(remote_file_name, local_file_name + ".tmp")
                
                
                # download succeeded: rename file.tmp -> file
                os.rename(local_file_name + ".tmp", local_file_name)
                cleanup_files.add(local_file_name)
                logger.info("Downloaded %s" % (local_file_name))
                return (action, local_file_name)
            except IOError, inst:
                error_msg = "%s\n" % (traceback.format_exc(1))
                if (tries == CNT_RETRIES - 1) or "No such file or directory" in error_msg:
                    error_messages += "%s\n" % (traceback.format_exc(1))
                    logger.debug("%s not found" % remote_file_name)
                    break

    logger.Error(error_messages)
    error_messages =error_messages.replace("\n",     "\n    ")
    raise Exception("UCSC Genomic sequence for %s could not be found.\n%s" % 
                    (ucsc_genome_name, error_messages))







        
        
        
        
#_________________________________________________________________________________________
#
#   2b) unzip_files
#
#_________________________________________________________________________________________
def __unzip_files(logger, action, local_file_name, unzipped_name):
    """
    Run gunzip or unzip
    """
    
    def basename (files):
        return tuple(os.path.basename(f) for f in files)
        

    file_actions = {"tar"   : "tar -xzOf %s > %s",
                    "gunzip": "gunzip -c %s > %s",
                    "unzip" : "unzip  -p %s > %s"}

    
    command_str = file_actions[action]

    # 
    #   has file been unzipped previously?
    # 
    if not os.path.exists(unzipped_name):
        
        message = "%s %s > %s" % basename((action, local_file_name, unzipped_name))
        logger.debug(message)
        result = os.system(command_str % (local_file_name, unzipped_name))
        if result != 0:
            raise Exception(message + " returned %s" % str(result))
        logger.info(message)
    else:
        logger.debug("%s already unzipped" % unzipped_name)
        
    
            
#_________________________________________________________________________________________
#
#   3) one_line_per_contig
#
#_________________________________________________________________________________________
def __one_line_per_contig(logger_, unzipped_name, fasta_file_name):
    """
    Convert fasta file to one line per sequence
    """
    if os.path.exists(fasta_file_name):
        logger_.debug("%s already one line per contig." % fasta_file_name)
        return

    tmp_fasta_path_strm   = open(unzipped_name)
    final_fasta_path_strm = open(fasta_file_name + ".tmp", "w")
    
    logger_.debug("Converting [%s] to single line format" % unzipped_name)
    (cnt_lines, cnt_sequences) = fasta_one_line_per_sequence (  tmp_fasta_path_strm, 
                                                                final_fasta_path_strm,
                                                                logger = logger_)        
    logger_.info("%s converted to single line format" % fasta_file_name)
    logger_.info("%d lines --> 1 line per %d contigs" % (cnt_lines, cnt_sequences))
    final_fasta_path_strm.close()
    
    # Use separate rename step in case processes interrupted
    #       and we have  partially created file
    os.rename(fasta_file_name + ".tmp", fasta_file_name)
    
        
            
            
    

#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Logger


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

import unittest
import logging

if __name__ == "__main__":
    logger = logging.getLogger(module_name)
    setup_std_logging(logger, options.log_file, options.verbose)

    if not options.skip_parameter_logging:
        programme_name = os.path.split(sys.argv[0])[1]
        logger.info("%s %s" % (programme_name, " ".join(get_option_strings(parser, options))))
    
#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888

#   Main logic


#88888888888888888888888888888888888888888888888888888888888888888888888888888888888888888


            

#
#   debug code not run if called as a module
#     
if __name__ == '__main__':
    if "--debug" in sys.argv:
        sys.argv.remove("--debug")
    sys.argv.append("--verbose")
    download_and_index_ucsc_genome_sequence(logger, 
                                            options.species_name,
                                            options.ucsc_genome_name,
                                            options.output_dir,
                                            options.goldenpath_ftp,
                                            options.temp_directory,
                                            options.verbose)

    
    

